rm(list = ls())

library(data.table)

load('./data/panel_month_dummies.RData')

permute <- function(permutation_data, n_reps) {
  results <- vector(mode='list', length=n_reps)
  for (r in seq(n_reps)) {
    #print(paste("Rep", r))
    true_treatments <- data.table(permutation_data)[, .(is_treated=unique(losing_income_d)==1), by=panelvar]
    
    permutation_treatments <- data.table(
      is_treated_permutation = sample(true_treatments$is_treated),
      permutation_id = unique(permutation_data$panelvar)
    )
    
    permutation_data$is_treated_permutation <- NULL
    permutation_data <- merge(permutation_data, permutation_treatments, by.x='panelvar', by.y='permutation_id')
    
    nb_b_c <- MASS::glm.nb(asian_hc ~ 1, data=permutation_data[covid2==0 & is_treated_permutation==0])
    nb_b_t <- MASS::glm.nb(asian_hc ~ 1, data=permutation_data[covid2==0 & is_treated_permutation==1])
    nb_a_c <- MASS::glm.nb(asian_hc ~ 1, data=permutation_data[covid2==1 & is_treated_permutation==0])
    nb_a_t <- MASS::glm.nb(asian_hc ~ 1, data=permutation_data[covid2==1 & is_treated_permutation==1])
    
    atet <- mean((predict(nb_a_t, permutation_data[covid2==1 & is_treated_permutation==1], type = 'response') -
                    predict(nb_b_t, permutation_data[covid2==1 & is_treated_permutation==1], type = 'response')) -
                   (predict(nb_a_c, permutation_data[covid2==1 & is_treated_permutation==1], type = 'response') -
                      predict(nb_b_c, permutation_data[covid2==1 & is_treated_permutation==1], type = 'response')))
    results[[r]] <- atet
  }
  return(results)
}


set.seed(879)
rr <- permute(panel, 500)

save(rr, file='./output/neg_binomial_ri.RData')